from llama_index.postprocessor.dashscope_rerank import DashScopeRerank
from llama_index.core.postprocessor import SimilarityPostprocessor
from langchain_community.chat_models import ChatTongyi
from llama_index.core.indices.query.query_transform.base import HyDEQueryTransform
from llama_index.core import Settings,SimpleDirectoryReader,VectorStoreIndex

from llama_index.embeddings.dashscope import (
    DashScopeEmbedding,
    DashScopeTextEmbeddingModels,
    DashScopeTextEmbeddingType
)
#重排序
#词嵌入模型
embed_model = DashScopeEmbedding(
    model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V3,
    text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
    api_key="sk-f97e3654139742a4b01a99631628d36d"
)

# 初始化LLM
llm = ChatTongyi(model="qwen-plus", api_key="sk-f97e3654139742a4b01a99631628d36d")

Settings.llm = llm
Settings.embed_model = embed_model

hyde = HyDEQueryTransform(include_original=True)

docs = SimpleDirectoryReader("D:\Code\sshcode\RAG_pro\docs").load_data()
index = VectorStoreIndex.from_documents(docs)
rerank_query_engine = index.as_query_engine(
    similarity_top_k = 4,
    streaming = True,
    node_postprocessors=[
        #设置重排序模型
        DashScopeRerank(top_n=2,model="gte-rerank-v2",api_key="sk-f97e3654139742a4b01a99631628d36d"),
        SimilarityPostprocessor(similarity_cutoff=0.2) #相似度小于0.2的就过滤
    ]
)

rerank_res = rerank_query_engine.query("张华在哪个部门?")
print(rerank_res.print_response_stream())